use crate::config::CONFIG;

use log::*;
// 引入随机数处理库
// use rand::prelude::*;
// 引入序列化与反序列化库
use serde::{Deserialize, Serialize};
// 引入 SQL 语法解析库
// use sqlparser::parser::{Parser, ParserError};
// use sqlparser::{ast::*, dialect::PostgreSqlDialect};
// 引入文件操作库
use std::fs::File;
use std::{fs, io::prelude::*};
// 引入异步 PostgreSQL 客户端库
// use crate::logger;
// use sqlparser::ast::Statement;
// use sqlparser::ast::{Ident, ObjectName};
// use sqlparser::tokenizer::{Token, Tokenizer};
use tokio_postgres::{Error, NoTls, Row};

use regex::Regex;
// use serde_json::{from_str, to_string_pretty};
// use std::io::prelude::*;
// use walkdir::WalkDir;

use rand::distributions::Uniform;
use rand::seq::SliceRandom;
use rand::Rng;
use std::collections::HashSet;

// 定义列信息结构体
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ColumnInfo {
    column_name: String,
    data_type: String,
    min: Option<String>,
    max: Option<String>,
    unique_values: Option<Vec<String>>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Relationship {
    column_name: String,
    foreign_table_name: String,
    foreign_column_name: String,
}
impl PartialEq for Relationship {
    fn eq(&self, other: &Self) -> bool {
        self.column_name == other.column_name
            && self.foreign_table_name == other.foreign_table_name
            && self.foreign_column_name == other.foreign_column_name
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct RelationshipHavetable {
    table: String,
    column_name: String,
    foreign_table_name: String,
    foreign_column_name: String,
}

// 定义表信息结构体
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct TableInfo {
    table_name: String,
    columns: Vec<ColumnInfo>,
    relationships: Option<Vec<Relationship>>,
}

// 异步获取表信息列表
pub async fn get_table_info_list(db_connect: &str) -> Result<Vec<TableInfo>, Error> {
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    // 连接到数据库
    let (client, connection) = tokio_postgres::connect(db_connect, NoTls).await?;
    // 在单独的任务中处理连接
    tokio::spawn(async move {
        if let Err(e) = connection.await {
            error!("Connection error: {}", e);
        }
    });

    // 获取表名
    let table_names = client
        .query(
            "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'",
            &[],
        )
        .await?;
    let mut table_info_list = Vec::new();
    info!("table_names {:?}", table_names);

    // 遍历表名，获取表信息
    for row in table_names {
        let table_name: &str = row.get(0);
        let columns = client
            .query(
                "SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = 'public' AND table_name = $1",
                &[&table_name],
            )
            .await?;

        let mut column_info_list = Vec::new();

        for col in columns {
            let column_name: &str = col.get(0);
            let data_type: &str = col.get(1);
            let is_numeric = [
                "integer",
                "bigint",
                "smallint",
                "numeric",
                "double precision",
                "real",
            ]
            .contains(&data_type);

            let (min, max, unique_values) = if is_numeric {
                let min_max: Row = client
                    .query_one(
                        format!(
                            "SELECT MIN(\"{}\") AS min, MAX(\"{}\") AS max FROM \"{}\"",
                            column_name, column_name, table_name
                        )
                        .as_str(),
                        &[],
                    )
                    .await?;

                let min = match data_type {
                    "integer" => min_max.get::<_, Option<i32>>(0).map(|v| v.to_string()),
                    "bigint" => min_max.get::<_, Option<i64>>(0).map(|v| v.to_string()),
                    "smallint" => min_max.get::<_, Option<i16>>(0).map(|v| v.to_string()),
                    "numeric" | "double precision" | "real" => {
                        min_max.get::<_, Option<f64>>(0).map(|v| v.to_string())
                    }
                    _ => None,
                };

                let max = match data_type {
                    "integer" => min_max.get::<_, Option<i32>>(1).map(|v| v.to_string()),
                    "bigint" => min_max.get::<_, Option<i64>>(1).map(|v| v.to_string()),
                    "smallint" => min_max.get::<_, Option<i16>>(1).map(|v| v.to_string()),
                    "numeric" | "double precision" | "real" => {
                        min_max.get::<_, Option<f64>>(1).map(|v| v.to_string())
                    }
                    _ => None,
                };

                (min, max, None)
            } else {
                let unique_values_result = client
                    .query(
                        format!(
                            "SELECT DISTINCT \"{}\" FROM \"{}\" ORDER BY \"{}\" LIMIT {}",
                            column_name, table_name, column_name, config.num_of_unique_values
                        )
                        .as_str(),
                        &[],
                    )
                    .await?;

                (
                    None,
                    None,
                    Some(
                        unique_values_result
                            .into_iter()
                            .filter_map(|row| row.get::<_, Option<String>>(0))
                            .collect::<Vec<String>>(),
                    ),
                )
            };

            column_info_list.push(ColumnInfo {
                column_name: column_name.to_string(),
                data_type: data_type.to_string(),
                min,
                max,
                unique_values,
            });
        }

        table_info_list.push(TableInfo {
            table_name: table_name.to_string(),
            columns: column_info_list,
            relationships: None,
        });
    }

    Ok(table_info_list)
}

// 将数据写入文件的函数
pub fn write_to_file(file_name: &str, data: &str) -> std::io::Result<()> {
    // 创建文件
    let mut file = File::create(file_name)?;
    // 将数据写入文件
    file.write_all(data.as_bytes())?;
    // 返回操作结果
    Ok(())
}

// 从文件中加载表信息列表
pub fn load_table_info_list_from_file(
    file_name: &str,
) -> Result<Vec<TableInfo>, Box<dyn std::error::Error>> {
    // 读取文件内容
    let contents = fs::read_to_string(file_name)?;
    // 将文件内容反序列化成TableInfo结构体的向量
    let table_info_list: Vec<TableInfo> = serde_json::from_str(&contents)?;
    // 返回表信息列表
    Ok(table_info_list)
}

// 为 table_info 增加外键依赖，保存在新文件中
pub fn add_foreign_dependent(sql_files_path: &str) {
    // let mut tables: Vec<TableInfo> =
    //     serde_json::from_str(include_str!("../table_info.json")).unwrap();
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    let table_info_contents =
        fs::read_to_string(&config.table_info_filename).expect("无法读取 table_info.json 文件");
    let mut tables: Vec<TableInfo> = serde_json::from_str(&table_info_contents).unwrap();

    let dir = sql_files_path;

    match fs::read_dir(dir) {
        Ok(entries) => {
            for entry in entries {
                let entry = entry.unwrap();
                let sql = fs::read_to_string(entry.path()).expect("Unable to read the file");

                let relationships = extract_relationships_from_query(&sql);
                for relationship in relationships {
                    add_relationship(
                        &mut tables,
                        &relationship.table,
                        &relationship.column_name,
                        &relationship.foreign_table_name,
                        &relationship.foreign_column_name,
                    );
                }
            }
        }
        Err(e) => {
            eprintln!(
                "Warning: Unable to read the directory: {}. 将生成没有外键关系约束的随机sql.",
                e
            );
        }
    }

    let updated_tables_json =
        serde_json::to_string_pretty(&tables).expect("Unable to convert to JSON");
    fs::write(&config.table_info_foreign_filename, updated_tables_json)
        .expect("Unable to write the file");
}

// 从查询中学到的外键依赖，由 add_foreign_dependent 函数调用
pub fn extract_relationships_from_query(query: &str) -> Vec<RelationshipHavetable> {
    let mut relationships: Vec<RelationshipHavetable> = Vec::new();

    let alias_regex = Regex::new(r"(?i)(?:FROM|,)\s+(?P<table>\w+)\s+as\s+(?P<alias>\w+)").unwrap();

    let mut alias_mapping = std::collections::HashMap::new();

    for cap in alias_regex.captures_iter(query) {
        let table = cap.name("table").unwrap().as_str().to_string();
        let alias = cap.name("alias").unwrap().as_str().to_string();
        alias_mapping.insert(alias, table);
    }

    debug!("alias_mapping {:?} ", alias_mapping);

    let regex = Regex::new(
        r"(?P<table1>[a-z]+\d*)\.(?P<column1>[a-z_]+)\s*=\s*(?P<table2>[a-z]+\d*)\.(?P<column2>[a-z_]+)",
    )
    .unwrap();
    for cap in regex.captures_iter(query) {
        let table1_alias = &cap["table1"];
        let column1 = &cap["column1"];
        let table2_alias = &cap["table2"];
        let column2 = &cap["column2"];

        let table1 = alias_mapping
            .get(table1_alias)
            .map(|s| s.as_str())
            .unwrap_or(table1_alias);
        let table2 = alias_mapping
            .get(table2_alias)
            .map(|s| s.as_str())
            .unwrap_or(table2_alias);

        debug!(
            "table1 {:?} table1_alias {:?} column1 {:?} table2 {:?} table2_alias {:?} column2 {:?}",
            table1, table1_alias, column1, table2, table2_alias, column2
        );

        relationships.push(RelationshipHavetable {
            table: table1.to_string(),
            column_name: column1.to_string(),
            foreign_table_name: table2.to_string(),
            foreign_column_name: column2.to_string(),
        });

        relationships.push(RelationshipHavetable {
            table: table2.to_string(),
            column_name: column2.to_string(),
            foreign_table_name: table1.to_string(),
            foreign_column_name: column1.to_string(),
        });
    }

    relationships
}

// 为json结构增加关系，由 add_foreign_dependent 函数调用
fn add_relationship(
    tables: &mut Vec<TableInfo>,
    org_table: &str,
    column_name: &str,
    foreign_table_name: &str,
    foreign_column_name: &str,
) {
    if let Some(table) = tables.iter_mut().find(|t| t.table_name == org_table) {
        if table.columns.iter().any(|c| c.column_name == column_name) {
            let relationship = Relationship {
                column_name: column_name.to_string(),
                foreign_table_name: foreign_table_name.to_string(),
                foreign_column_name: foreign_column_name.to_string(),
            };

            debug!("relationship {:?}", relationship);

            // Check if the relationship already exists in the relationships vector
            let relationship_exists = if let Some(relationships) = &table.relationships {
                relationships.iter().any(|r| r == &relationship)
            } else {
                false
            };

            if !relationship_exists {
                if let Some(relationships) = &mut table.relationships {
                    relationships.push(relationship);
                } else {
                    table.relationships = Some(vec![relationship]);
                }
            }
        }
    }
}

// 生成新的FROM子句
fn generate_from_clause(tables: &Vec<TableInfo>) -> String {
    tables
        .iter()
        .map(|table| table.table_name.clone())
        .collect::<Vec<String>>()
        .join(", ")
}

// 生成连接条件
fn generate_join_conditions(selected_tables: &[TableInfo]) -> HashSet<String> {
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    let mut rng = rand::thread_rng();
    let mut join_conditions = HashSet::new();
    let probability: f64 = config.join_conditions_probability;

    for table in selected_tables.iter() {
        if let Some(ref relationships) = table.relationships {
            for rel in relationships {
                if selected_tables
                    .iter()
                    .any(|t| t.table_name == rel.foreign_table_name)
                {
                    let should_add_condition = rng.gen_range(0.0..=1.0) < probability;
                    if should_add_condition {
                        let mut condition_parts = vec![
                            format!("{}.{}", table.table_name, rel.column_name),
                            format!("{}.{}", rel.foreign_table_name, rel.foreign_column_name),
                        ];
                        condition_parts.sort();
                        let condition = condition_parts.join(" = ");
                        join_conditions.insert(condition);
                    }
                }
            }
        }
    }

    join_conditions
}

// 生成约束条件
fn generate_constraints(selected_tables: &[TableInfo]) -> Vec<String> {
    let mut rng = rand::thread_rng();
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    // 随机生成约束条件
    // 定义约束条件数量的随机范围 // have done 本逻辑中大量随机超参数，可以定义到 postgres.cfg 中
    let num_constraints_range = Uniform::from(0..=config.num_constraints_range);
    // 随机生成约束条件数量
    let num_constraints = rng.sample(num_constraints_range);
    // 随机生成约束条件
    let mut constraints = Vec::new();

    // 添加一个权重，使得整数和浮点数类型的列具有更高的优先级
    // fn column_weight(column: &ColumnInfo) -> usize {
    //     match column.data_type.as_str() {
    //         "integer" | "float" => 200, // 为整数和浮点数类型的列分配更高的权重
    //         _ => 1,
    //     }
    // }

    for table in selected_tables.iter() {
        for column in table.columns.iter() {
            // 权重修改
            // let weight = column_weight(column);
            // let selected = rng.gen_range(0..weight) == 0;

            // if selected {
            //     continue;
            // }

            // 判断 unique_values 是否存在且不为空
            let has_unique_values = column
                .unique_values
                .as_ref()
                .map_or(false, |v| !v.is_empty());

            if has_unique_values {
                let ref unique_values = column.unique_values.as_ref().unwrap();
                // 处理 unique_values 存在且不为空的情况
                let constraint_probability = rng.gen_range(0.0..=1.0);
                let constraint = if constraint_probability < config.constraint_probability {
                    let num_values = rng.gen_range(1..=config.num_values);
                    let selected_values = unique_values
                        .choose_multiple(&mut rng, num_values)
                        .cloned()
                        .map(|v| format!("$${}$$", v))
                        .collect::<Vec<String>>()
                        .join(", ");
                    format!(
                        "({}.{} in ({}))",
                        table.table_name, column.column_name, selected_values
                    )
                } else {
                    let null_probability = rng.gen_range(0.0..=1.0);
                    if null_probability < config.null_probability {
                        format!("({}.{} is not NULL)", table.table_name, column.column_name)
                    } else {
                        format!("({}.{} is NULL)", table.table_name, column.column_name)
                    }
                };
                constraints.push(constraint);
            } else {
                // unique_values 为空或不存在的情况

                // 如果是整数或浮点数类型，尝试使用 min 和 max 值
                if column.data_type == "integer" || column.data_type == "float" {
                    // 生成基于范围的约束条件

                    // 这里处理 column.unique_values 为 None 的情况

                    debug!("column.data_type == integer || column.data_type == float");
                    debug!(
                        "table {:?} column {:?} min_value {:?} max_value {:?}",
                        table, column, &column.min, &column.max
                    );
                    debug!("column.column_name {:?}", column.column_name);
                    if let (Some(min_value), Some(max_value)) = (&column.min, &column.max) {
                        let random_value = if column.data_type == "integer" {
                            rng.gen_range(
                                min_value.parse::<i32>().unwrap()
                                    ..=max_value.parse::<i32>().unwrap(),
                            )
                            .to_string()
                        } else {
                            rng.gen_range(
                                min_value.parse::<f32>().unwrap()
                                    ..=max_value.parse::<f32>().unwrap(),
                            )
                            .to_string()
                        };

                        let operators = vec![">", "<", "<=", ">=", "="];
                        let random_operator = operators.choose(&mut rng).unwrap();
                        let constraint = format!(
                            "({}.{} {} {})",
                            table.table_name, column.column_name, random_operator, random_value
                        );
                        debug!("{:?}", constraint);
                        constraints.push(constraint);
                    }
                }
            }
        }
    }

    // 在约束条件中随机选择一定数量的条件
    let constraints = constraints
        .choose_multiple(&mut rng, num_constraints)
        .cloned()
        .collect::<Vec<String>>();

    constraints
}

// 生成LIMIT子句
fn generate_limit_clause() -> String {
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    let mut rng = rand::thread_rng();
    let limit_range_probability = rng.gen_range(0.0..=1.0);

    let mut limit_value = config.limit_ranges.last().unwrap().1;
    for (probability, value) in &config.limit_ranges {
        // 这里已经使用了引用
        if limit_range_probability < *probability {
            // 使用 * 来解引用
            limit_value = rng.gen_range(1..=*value); // 使用 * 来解引用
            break;
        }
    }
    format!("LIMIT {}", limit_value)
}

// fn generate_limit_clause() -> String {
//     let mut rng = rand::thread_rng();
//     let limit_range_probability = rng.gen_range(0.0..=1.0);
//     let limit_value = if limit_range_probability < 0.25 {
//         rng.gen_range(1..=10)
//     } else if limit_range_probability < 0.5 {
//         rng.gen_range(10..=100)
//     } else if limit_range_probability < 0.75 {
//         rng.gen_range(100..=1000)
//     } else {
//         rng.gen_range(1000..=10000)
//     };

//     format!("LIMIT {}", limit_value)
// }

// 生成SELECT子句
fn generate_select_fields(selected_tables: &[TableInfo]) -> String {
    let mut rng = rand::thread_rng();
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    // let aggregate_functions = vec!["MIN", "MAX", "AVG", "SUM"];
    let aggregate_functions = vec!["MIN", "MAX"];

    let select_fields_probability = rng.gen_range(0.0..=1.0);
    let select_fields = if select_fields_probability < config.select_fields_probability {
        "COUNT(*)".to_string()
    } else {
        let num_columns = rng.gen_range(1..=config.num_max_minormax);
        let columns = selected_tables
            .iter()
            .flat_map(|t| {
                t.columns
                    .iter()
                    .map(|c| (t.table_name.clone(), c.column_name.clone()))
            })
            .collect::<Vec<(String, String)>>()
            .choose_multiple(&mut rng, num_columns)
            .map(|(table_name, column_name)| {
                // 随机选择一个聚合函数
                let random_agg_function = aggregate_functions.choose(&mut rng).unwrap();
                format!("{}({}.{})", random_agg_function, table_name, column_name)
            })
            .collect::<Vec<String>>()
            .join(", ");
        columns
    };

    select_fields
}

// 定义一个名为generate_random_sql的公共函数，参数为一个&Vec<TableInfo>
pub fn generate_random_sql(metadata: &Vec<TableInfo>) -> String {
    // 读取全局配置
    let config = CONFIG.read().unwrap();

    let mut rng = rand::thread_rng();
    let newline = if cfg!(windows) { "\r\n" } else { "\n" };

    // SQL中必要的部分
    let select_clause = "SELECT"; // 这里可以根据实际需求进行更改

    let where_clause = "WHERE";

    // 随机选取一定数量的表
    let num_tables = rng.gen_range(1..=config.num_max_tables);
    let selected_tables: Vec<TableInfo> = metadata
        .choose_multiple(&mut rng, num_tables)
        .cloned()
        .collect();

    let from_clause = generate_from_clause(&selected_tables);
    let join_conditions = generate_join_conditions(&selected_tables);
    let constraints = generate_constraints(&selected_tables);
    let limit_clause = generate_limit_clause();
    let select_fields = generate_select_fields(&selected_tables);

    // 组合新的sql
    let conditions_string = join_conditions
        .into_iter()
        .chain(constraints)
        .collect::<Vec<String>>()
        .join(" AND ");
    let new_sql = if conditions_string.is_empty() {
        format!(
            "{} {}{}FROM {}{}{}",
            select_clause,
            select_fields,
            newline,
            from_clause.replace(", ", &format!(",{}", newline)), // 在表名之间添加换行符
            newline,
            limit_clause
        )
    } else {
        format!(
            "{} {}{}FROM {}{}{} {}{}{}",
            select_clause,
            select_fields,
            newline,
            from_clause.replace(", ", &format!(",{}", newline)), // 在表名之间添加换行符
            newline,
            where_clause,
            conditions_string.replace(" AND ", &format!(" {}AND ", newline)), // 在AND之前添加换行符
            newline,
            limit_clause
        )
    };

    new_sql
}
